package crypto

import (
	"bytes"
	"encoding/binary"
    "encoding/json"
	"fmt"
    "sort"
	"reflect"
    "math/big"
)

// Serializer provides serialization and deserialization for cryptographic types
type Serializer struct {
	// Type markers
	TypeMarkers map[reflect.Type]byte
	ReverseMarkers map[byte]reflect.Type
    // Mode toggles; when PyCompat is true, use Python-like markers for composites
    PyCompat bool
}

// Tuple is a lightweight container to preserve tuple framing in PyCompat
type Tuple struct { Items []interface{} }

// NewSerializer creates a new Serializer instance
func NewSerializer() *Serializer {
	s := &Serializer{
		TypeMarkers: make(map[reflect.Type]byte),
		ReverseMarkers: make(map[byte]reflect.Type),
        PyCompat: false,
	}
	
	// Initialize type markers
	s.TypeMarkers[reflect.TypeOf(G1{})] = 0x01
	s.TypeMarkers[reflect.TypeOf(G2{})] = 0x02
	s.TypeMarkers[reflect.TypeOf(ZR{})] = 0x03
	s.TypeMarkers[reflect.TypeOf(GT{})] = 0x04
	s.TypeMarkers[reflect.TypeOf([]byte{})] = 0x05
	s.TypeMarkers[reflect.TypeOf("")] = 0x06
	s.TypeMarkers[reflect.TypeOf(int(0))] = 0x07
	s.TypeMarkers[reflect.TypeOf(int64(0))] = 0x08
	s.TypeMarkers[reflect.TypeOf([]interface{}{})] = 0x09
	s.TypeMarkers[reflect.TypeOf(map[string]interface{}{})] = 0x0A
	
	// Build reverse mapping
	for t, marker := range s.TypeMarkers {
		s.ReverseMarkers[marker] = t
	}
	
	return s
}

// Serialize serializes an object to bytes
func (s *Serializer) Serialize(obj interface{}) ([]byte, error) {
	var buf bytes.Buffer
    
	if err := s.serializeValue(&buf, obj); err != nil {
		return nil, fmt.Errorf("serialization failed: %w", err)
	}
    
	return buf.Bytes(), nil
}

// Deserialize deserializes bytes to an object
func (s *Serializer) Deserialize(data []byte) (interface{}, error) {
	buf := bytes.NewBuffer(data)
	return s.deserializeValue(buf)
}

// serializeValue serializes a value to the buffer
func (s *Serializer) serializeValue(buf *bytes.Buffer, value interface{}) error {
	if value == nil {
		return s.serializeNil(buf)
	}
	
	v := reflect.ValueOf(value)
	t := v.Type()
	
	// Handle different types
	switch v.Kind() {
	case reflect.Slice:
		return s.serializeSlice(buf, v)
	case reflect.Map:
		return s.serializeMap(buf, v)
	case reflect.Struct:
		return s.serializeStruct(buf, v)
	case reflect.String:
		return s.serializeString(buf, v.String())
	case reflect.Int, reflect.Int64:
		return s.serializeInt(buf, v.Int())
	case reflect.Uint, reflect.Uint64:
		return s.serializeUint(buf, v.Uint())
	default:
		// Try to find a type marker
		if marker, ok := s.TypeMarkers[t]; ok {
			return s.serializeWithMarker(buf, marker, value)
		}
		return fmt.Errorf("unsupported type: %v", t)
	}
}

// serializeWithMarker serializes with a type marker
func (s *Serializer) serializeWithMarker(buf *bytes.Buffer, marker byte, value interface{}) error {
	// Write type marker
	if err := buf.WriteByte(marker); err != nil {
		return err
	}
	
	// Serialize based on type
	switch v := value.(type) {
	case G1:
		return s.serializeG1(buf, &v)
	case G2:
		return s.serializeG2(buf, &v)
	case ZR:
		return s.serializeZR(buf, &v)
	case GT:
		return s.serializeGT(buf, &v)
	case []byte:
		return s.serializeBytes(buf, v)
	case string:
		return s.serializeString(buf, v)
	case int:
		return s.serializeInt(buf, int64(v))
	case int64:
		return s.serializeInt(buf, v)
	default:
		return fmt.Errorf("unsupported type for marker %x: %T", marker, value)
	}
}

// serializeNil serializes nil value
func (s *Serializer) serializeNil(buf *bytes.Buffer) error {
	return buf.WriteByte(0x00)
}

// serializeSlice serializes a slice
func (s *Serializer) serializeSlice(buf *bytes.Buffer, v reflect.Value) error {
	// Write slice marker
	if err := buf.WriteByte(0x09); err != nil {
		return err
	}
	
	// Write length
	length := v.Len()
	if err := binary.Write(buf, binary.LittleEndian, uint32(length)); err != nil {
		return err
	}
	
	// Serialize elements
	for i := 0; i < length; i++ {
		if err := s.serializeValue(buf, v.Index(i).Interface()); err != nil {
			return err
		}
	}
	
	return nil
}

// serializeMap serializes a map
func (s *Serializer) serializeMap(buf *bytes.Buffer, v reflect.Value) error {
	// Write map marker
	if err := buf.WriteByte(0x0A); err != nil {
		return err
	}
	
	// Write length
	length := v.Len()
	if err := binary.Write(buf, binary.LittleEndian, uint32(length)); err != nil {
		return err
	}
	
	// Serialize key-value pairs
	for _, key := range v.MapKeys() {
		// Serialize key
		if err := s.serializeValue(buf, key.Interface()); err != nil {
			return err
		}
		// Serialize value
		if err := s.serializeValue(buf, v.MapIndex(key).Interface()); err != nil {
			return err
		}
	}
	
	return nil
}

// serializeStruct serializes a struct
func (s *Serializer) serializeStruct(buf *bytes.Buffer, v reflect.Value) error {
	// Write struct marker
	if err := buf.WriteByte(0x0B); err != nil {
		return err
	}
	
	// Write number of fields
	numFields := v.NumField()
	if err := binary.Write(buf, binary.LittleEndian, uint32(numFields)); err != nil {
		return err
	}
	
	// Serialize fields
	for i := 0; i < numFields; i++ {
		field := v.Field(i)
		if field.CanInterface() {
			if err := s.serializeValue(buf, field.Interface()); err != nil {
				return err
			}
		}
	}
	
	return nil
}

// serializeG1 serializes a G1 element
func (s *Serializer) serializeG1(buf *bytes.Buffer, g1 *G1) error {
    b := g1.G1Affine.Bytes()
    return s.serializeBytes(buf, b[:])
}

// serializeG2 serializes a G2 element
func (s *Serializer) serializeG2(buf *bytes.Buffer, g2 *G2) error {
    b := g2.G2Affine.Bytes()
    return s.serializeBytes(buf, b[:])
}

// serializeZR serializes a ZR element
func (s *Serializer) serializeZR(buf *bytes.Buffer, zr *ZR) error {
    b := zr.Element.Bytes()
    return s.serializeBytes(buf, b[:])
}

// serializeGT serializes a GT element
func (s *Serializer) serializeGT(buf *bytes.Buffer, gt *GT) error {
    b := gt.GT.Bytes()
    return s.serializeBytes(buf, b[:])
}

// serializeBytes serializes a byte slice
func (s *Serializer) serializeBytes(buf *bytes.Buffer, data []byte) error {
	// Write length
	if err := binary.Write(buf, binary.LittleEndian, uint32(len(data))); err != nil {
		return err
	}
	
	// Write data
	_, err := buf.Write(data)
	return err
}

// serializeString serializes a string
func (s *Serializer) serializeString(buf *bytes.Buffer, str string) error {
	data := []byte(str)
	return s.serializeBytes(buf, data)
}

// serializeInt serializes an integer
func (s *Serializer) serializeInt(buf *bytes.Buffer, val int64) error {
	return binary.Write(buf, binary.LittleEndian, val)
}

// serializeUint serializes an unsigned integer
func (s *Serializer) serializeUint(buf *bytes.Buffer, val uint64) error {
	return binary.Write(buf, binary.LittleEndian, val)
}

// deserializeValue deserializes a value from the buffer
func (s *Serializer) deserializeValue(buf *bytes.Buffer) (interface{}, error) {
	if buf.Len() == 0 {
		return nil, fmt.Errorf("unexpected end of data")
	}
    
	// Read type marker
	marker, err := buf.ReadByte()
	if err != nil {
		return nil, err
	}
	
	// Handle nil
	if marker == 0x00 {
		return nil, nil
	}
	
	// Deserialize based on marker
	switch marker {
	case 0x01:
		return s.deserializeG1(buf)
	case 0x02:
		return s.deserializeG2(buf)
	case 0x03:
		return s.deserializeZR(buf)
	case 0x04:
		return s.deserializeGT(buf)
	case 0x05:
		return s.deserializeBytes(buf)
	case 0x06:
		return s.deserializeString(buf)
	case 0x07, 0x08:
		return s.deserializeInt(buf)
	case 0x09:
		return s.deserializeSlice(buf)
	case 0x0A:
		return s.deserializeMap(buf)
	case 0x0B:
		return s.deserializeStruct(buf)
	default:
		return nil, fmt.Errorf("unknown type marker: %x", marker)
	}
}

// Python-compatible (HBACSS) codec

const (
    pyListStart  = byte(0x01)
    pyListEnd    = byte(0x02)
    pyTupleStart = byte(0x03)
    pyTupleEnd   = byte(0x04)
    pyG1Marker   = byte(0x05)
    pyG2Marker   = byte(0x06)
    pyZRMarker   = byte(0x07)
    pyGTMarker   = byte(0x08)
    pyBytes      = byte(0x09)
    pyMapStart   = byte(0x0A)
    pyMapEnd     = byte(0x0B)
    pyMisc       = byte(0x99)
)

// SerializePyCompat serializes using Python-like markers and framing for composites.
// For G1/G2/ZR, we write a length (u32 LE) + raw bytes to be tolerant to 48/96 sizes.
func (s *Serializer) SerializePyCompat(obj interface{}) ([]byte, error) {
    var buf bytes.Buffer
    s.PyCompat = true
    defer func() { s.PyCompat = false }()
    if err := s.pySerializeValue(&buf, obj); err != nil { return nil, err }
    return buf.Bytes(), nil
}

// DeserializePyCompat deserializes Python-like stream.
func (s *Serializer) DeserializePyCompat(data []byte) (interface{}, error) {
    s.PyCompat = true
    defer func() { s.PyCompat = false }()
    buf := bytes.NewBuffer(data)
    return s.pyDeserializeValue(buf)
}

func (s *Serializer) pySerializeValue(buf *bytes.Buffer, v interface{}) error {
    if v == nil { return buf.WriteByte(pyMisc) }
    switch val := v.(type) {
    case []interface{}:
        if err := buf.WriteByte(pyListStart); err != nil { return err }
        for _, it := range val {
            if err := s.pySerializeValue(buf, it); err != nil { return err }
        }
        return buf.WriteByte(pyListEnd)
    case map[string]interface{}:
        if err := buf.WriteByte(pyMapStart); err != nil { return err }
        // Write number of keys (deterministic order below)
        if err := binary.Write(buf, binary.LittleEndian, uint32(len(val))); err != nil { return err }
        // Deterministic key order: sort by key (lexicographic)
        keys := make([]string, 0, len(val))
        for k := range val { keys = append(keys, k) }
        sort.Strings(keys)
        for _, k := range keys {
            if err := s.pySerializeValue(buf, []byte(k)); err != nil { return err }
            if err := s.pySerializeValue(buf, val[k]); err != nil { return err }
        }
        return buf.WriteByte(pyMapEnd)
    case []byte:
        if err := buf.WriteByte(pyBytes); err != nil { return err }
        return s.serializeBytes(buf, val)
    case Tuple:
        if err := buf.WriteByte(pyTupleStart); err != nil { return err }
        for _, it := range val.Items {
            if err := s.pySerializeValue(buf, it); err != nil { return err }
        }
        return buf.WriteByte(pyTupleEnd)
    case G1:
        if err := buf.WriteByte(pyG1Marker); err != nil { return err }
        b := val.G1Affine.Bytes()
        return s.serializeBytes(buf, b[:])
    case *G1:
        if err := buf.WriteByte(pyG1Marker); err != nil { return err }
        b := val.G1Affine.Bytes()
        return s.serializeBytes(buf, b[:])
    case G2:
        if err := buf.WriteByte(pyG2Marker); err != nil { return err }
        b := val.G2Affine.Bytes()
        return s.serializeBytes(buf, b[:])
    case *G2:
        if err := buf.WriteByte(pyG2Marker); err != nil { return err }
        b := val.G2Affine.Bytes()
        return s.serializeBytes(buf, b[:])
    case ZR:
        if err := buf.WriteByte(pyZRMarker); err != nil { return err }
        b := val.Element.Bytes()
        return s.serializeBytes(buf, b[:])
    case *ZR:
        if err := buf.WriteByte(pyZRMarker); err != nil { return err }
        b := val.Element.Bytes()
        return s.serializeBytes(buf, b[:])
    case GT:
        if err := buf.WriteByte(pyGTMarker); err != nil { return err }
        b := val.GT.Bytes()
        return s.serializeBytes(buf, b[:])
    case *GT:
        if err := buf.WriteByte(pyGTMarker); err != nil { return err }
        b := val.GT.Bytes()
        return s.serializeBytes(buf, b[:])
    default:
        // Fallback: JSON-encode as misc
        if err := buf.WriteByte(pyMisc); err != nil { return err }
        jb, _ := json.Marshal(val)
        return s.serializeBytes(buf, jb)
    }
}

func (s *Serializer) pyDeserializeValue(buf *bytes.Buffer) (interface{}, error) {
    if buf.Len() == 0 { return nil, fmt.Errorf("unexpected eof") }
    marker, err := buf.ReadByte(); if err != nil { return nil, err }
    switch marker {
    case pyListStart:
        // read until list_end; collect items
        out := []interface{}{}
        for {
            if buf.Len() == 0 { return nil, fmt.Errorf("unclosed list") }
            // peek next
            b, _ := buf.ReadByte()
            if b == pyListEnd { break }
            buf.UnreadByte()
            it, err := s.pyDeserializeValue(buf)
            if err != nil { return nil, err }
            out = append(out, it)
        }
        return out, nil
    case pyMapStart:
        // Read length
        var ln uint32
        if err := binary.Read(buf, binary.LittleEndian, &ln); err != nil { return nil, err }
        m := make(map[string]interface{}, ln)
        for i := uint32(0); i < ln; i++ {
            kRaw, err := s.pyDeserializeValue(buf); if err != nil { return nil, err }
            // Keys are stored as raw bytes
            kb, ok := kRaw.([]byte); if !ok { return nil, fmt.Errorf("bad map key type") }
            val, err := s.pyDeserializeValue(buf); if err != nil { return nil, err }
            m[string(kb)] = val
        }
        // Read end marker
        end, err := buf.ReadByte(); if err != nil { return nil, err }
        if end != pyMapEnd { return nil, fmt.Errorf("map not properly terminated") }
        return m, nil
    case pyTupleStart:
        // Read until tuple_end
        items := []interface{}{}
        for {
            if buf.Len() == 0 { return nil, fmt.Errorf("unclosed tuple") }
            b, _ := buf.ReadByte()
            if b == pyTupleEnd { break }
            buf.UnreadByte()
            it, err := s.pyDeserializeValue(buf)
            if err != nil { return nil, err }
            items = append(items, it)
        }
        return Tuple{Items: items}, nil
    case pyBytes:
        return s.deserializeBytes(buf)
    case pyG1Marker:
        b, err := s.deserializeBytes(buf); if err != nil { return nil, err }
        var g G1
        if _, err := g.G1Affine.SetBytes(b); err != nil { return nil, err }
        return g, nil
    case pyG2Marker:
        b, err := s.deserializeBytes(buf); if err != nil { return nil, err }
        var g G2
        if _, err := g.G2Affine.SetBytes(b); err != nil { return nil, err }
        return g, nil
    case pyZRMarker:
        b, err := s.deserializeBytes(buf); if err != nil { return nil, err }
        var z ZR
        z.Element.SetBytes(b)
        return z, nil
    case pyGTMarker:
        b, err := s.deserializeBytes(buf); if err != nil { return nil, err }
        var t GT
        if err := t.GT.SetBytes(b); err != nil { return nil, err }
        return t, nil
    case pyListEnd:
        // should be handled by list parser
        return nil, fmt.Errorf("unexpected list_end")
    case pyTupleEnd:
        return nil, fmt.Errorf("unexpected bare tuple marker")
    case pyMisc:
        b, err := s.deserializeBytes(buf); if err != nil { return nil, err }
        var any interface{}
        if err := json.Unmarshal(b, &any); err != nil { return b, nil }
        return any, nil
    default:
        return nil, fmt.Errorf("unknown compat marker: 0x%x", marker)
    }
}

// deserializeG1 deserializes a G1 element
func (s *Serializer) deserializeG1(buf *bytes.Buffer) (*G1, error) {
    b, err := s.deserializeBytes(buf)
    if err != nil {
        return nil, err
    }
    var g1 G1
    if _, err := g1.G1Affine.SetBytes(b); err != nil { return nil, err }
    return &g1, nil
}

// deserializeG2 deserializes a G2 element
func (s *Serializer) deserializeG2(buf *bytes.Buffer) (*G2, error) {
    b, err := s.deserializeBytes(buf)
    if err != nil {
        return nil, err
    }
    var g2 G2
    if _, err := g2.G2Affine.SetBytes(b); err != nil { return nil, err }
    return &g2, nil
}

// deserializeZR deserializes a ZR element
func (s *Serializer) deserializeZR(buf *bytes.Buffer) (*ZR, error) {
    d, err := s.deserializeBytes(buf)
    if err != nil {
        return nil, err
    }
    
    var zr ZR
    zr.SetBigInt(new(big.Int).SetBytes(d))
	
	return &zr, nil
}

// deserializeGT deserializes a GT element
func (s *Serializer) deserializeGT(buf *bytes.Buffer) (*GT, error) {
    b, err := s.deserializeBytes(buf)
    if err != nil {
        return nil, err
    }
    var t GT
    if err := t.GT.SetBytes(b); err != nil { return nil, err }
    return &t, nil
}

// deserializeBytes deserializes a byte slice
func (s *Serializer) deserializeBytes(buf *bytes.Buffer) ([]byte, error) {
	// Read length
	var length uint32
	if err := binary.Read(buf, binary.LittleEndian, &length); err != nil {
		return nil, err
	}
	
	// Read data
	data := make([]byte, length)
	if _, err := buf.Read(data); err != nil {
		return nil, err
	}
	
	return data, nil
}

// deserializeString deserializes a string
func (s *Serializer) deserializeString(buf *bytes.Buffer) (string, error) {
	data, err := s.deserializeBytes(buf)
	if err != nil {
		return "", err
	}
	
	return string(data), nil
}

// deserializeInt deserializes an integer
func (s *Serializer) deserializeInt(buf *bytes.Buffer) (int64, error) {
	var val int64
	err := binary.Read(buf, binary.LittleEndian, &val)
	return val, err
}

// deserializeSlice deserializes a slice
func (s *Serializer) deserializeSlice(buf *bytes.Buffer) ([]interface{}, error) {
	// Read length
	var length uint32
	if err := binary.Read(buf, binary.LittleEndian, &length); err != nil {
		return nil, err
	}
	
	// Deserialize elements
	result := make([]interface{}, length)
	for i := uint32(0); i < length; i++ {
		val, err := s.deserializeValue(buf)
		if err != nil {
			return nil, err
		}
		result[i] = val
	}
	
	return result, nil
}

// deserializeMap deserializes a map
func (s *Serializer) deserializeMap(buf *bytes.Buffer) (map[string]interface{}, error) {
	// Read length
	var length uint32
	if err := binary.Read(buf, binary.LittleEndian, &length); err != nil {
		return nil, err
	}
	
	// Deserialize key-value pairs
	result := make(map[string]interface{})
	for i := uint32(0); i < length; i++ {
		// Deserialize key
		keyVal, err := s.deserializeValue(buf)
		if err != nil {
			return nil, err
		}
		
		// Deserialize value
		val, err := s.deserializeValue(buf)
		if err != nil {
			return nil, err
		}
		
		// Convert key to string
		keyStr := fmt.Sprintf("%v", keyVal)
		result[keyStr] = val
	}
	
	return result, nil
}

// deserializeStruct deserializes a struct
func (s *Serializer) deserializeStruct(buf *bytes.Buffer) (map[string]interface{}, error) {
	// Read number of fields
	var numFields uint32
	if err := binary.Read(buf, binary.LittleEndian, &numFields); err != nil {
		return nil, err
	}
	
	// Deserialize fields
	result := make(map[string]interface{})
	for i := uint32(0); i < numFields; i++ {
		val, err := s.deserializeValue(buf)
		if err != nil {
			return nil, err
		}
		result[fmt.Sprintf("field_%d", i)] = val
	}
	
	return result, nil
}

// Global serializer instance
var DefaultSerializer = NewSerializer()

// Convenience functions using the global instance

// Serialize serializes an object using the default serializer
func Serialize(obj interface{}) ([]byte, error) {
	return DefaultSerializer.Serialize(obj)
}

// Deserialize deserializes bytes using the default serializer
func Deserialize(data []byte) (interface{}, error) {
	return DefaultSerializer.Deserialize(data)
}

// SerializeG1 serializes a G1 element
func SerializeG1(g1 *G1) []byte {
    b := g1.G1Affine.Bytes()
    return b[:]
}

// DeserializeG1 deserializes bytes to a G1 element
func DeserializeG1(data []byte) (*G1, error) {
	return DefaultSerializer.deserializeG1(bytes.NewBuffer(data))
}

// SerializeG2 serializes a G2 element
func SerializeG2(g2 *G2) []byte {
    b := g2.G2Affine.Bytes()
    return b[:]
}

// DeserializeG2 deserializes bytes to a G2 element
func DeserializeG2(data []byte) (*G2, error) {
	return DefaultSerializer.deserializeG2(bytes.NewBuffer(data))
}

// SerializeGT serializes a GT element
func SerializeGT(gt *GT) []byte {
    b := gt.GT.Bytes()
    return b[:]
}

// DeserializeGT deserializes bytes to a GT element
func DeserializeGT(data []byte) (*GT, error) {
    buf := bytes.NewBuffer(nil)
    // Reuse same framing as other elements: write length then data for local helper
    // But here, we can directly set since we know input is raw bytes of GT
    var t GT
    if err := t.GT.SetBytes(data); err != nil { return nil, err }
    _ = buf // placeholder to avoid unused warning if build tags change
    return &t, nil
}
